﻿@page "/chat"
@using System.ComponentModel
@using AbpAiChat.Services
@using Volo.Abp.Users
@inject IChatClient ChatClient
@inject NavigationManager Nav
@inject SemanticSearch Search
@implements IDisposable

<PageTitle>Chat</PageTitle>

<ChatHeader OnNewChat="@ResetConversationAsync" />

<ChatMessageList Messages="@messages" InProgressMessage="@currentResponseMessage">
    <NoMessagesContent>
        <div>To get started, try asking about these example documents. You can replace these with your own data and replace this message.</div>
        <ChatCitation File="Example_Emergency_Survival_Kit.pdf"/>
        <ChatCitation File="Example_GPS_Watch.pdf"/>
    </NoMessagesContent>
</ChatMessageList>

<div class="chat-container">
    <ChatSuggestions OnSelected="@AddUserMessageAsync" @ref="@chatSuggestions" />
    <ChatInput OnSend="@AddUserMessageAsync" @ref="@chatInput" />
</div>

@code {
    private const string SystemPrompt = @"
        You are an assistant who answers questions about information you retrieve.
        Do not answer questions about anything else.
        Use only simple markdown to format your responses.

        Use the search tool to find relevant information. When you do this, end your
        reply with citations in the special XML format:

        <citation filename='string' page_number='number'>exact quote here</citation>

        Always include the citation in your response if there are results.

        The quote must be max 5 words, taken word-for-word from the search result, and is the basis for why the citation is relevant.
        Don't refer to the presence of citations; just emit these tags right at the end, with no surrounding text.
        ";

    private readonly ChatOptions chatOptions = new();
    private readonly List<ChatMessage> messages = new();
    private CancellationTokenSource? currentResponseCancellation;
    private ChatMessage? currentResponseMessage;
    private ChatInput? chatInput;
    private ChatSuggestions? chatSuggestions;

    [Inject]
    public ICurrentUser CurrentUser { get; set; }

    protected override void OnInitialized()
    {
        messages.Add(new(ChatRole.System, SystemPrompt));
        chatOptions.Tools =
        [
            AIFunctionFactory.Create(SearchAsync),
            AIFunctionFactory.Create(GetWeather),
            AIFunctionFactory.Create(GetCurrentUserInfo)
        ];
    }

    private async Task AddUserMessageAsync(ChatMessage userMessage)
    {
        CancelAnyCurrentResponse();

        // Add the user message to the conversation
        messages.Add(userMessage);
        chatSuggestions?.Clear();
        await chatInput!.FocusAsync();

        // Stream and display a new response from the IChatClient
        var responseText = new TextContent("");
        currentResponseMessage = new ChatMessage(ChatRole.Assistant, [responseText]);
        currentResponseCancellation = new();
        await foreach (var update in ChatClient.GetStreamingResponseAsync([.. messages], chatOptions, currentResponseCancellation.Token))
        {
            messages.AddMessages(update, filter: c => c is not TextContent);
            responseText.Text += update.Text;
            ChatMessageItem.NotifyChanged(currentResponseMessage);
        }

        // Store the final response in the conversation, and begin getting suggestions
        messages.Add(currentResponseMessage!);
        currentResponseMessage = null;
        chatSuggestions?.Update(messages);
    }

    private void CancelAnyCurrentResponse()
    {
        // If a response was cancelled while streaming, include it in the conversation so it's not lost
        if (currentResponseMessage is not null)
        {
            messages.Add(currentResponseMessage);
        }

        currentResponseCancellation?.Cancel();
        currentResponseMessage = null;
    }

    private async Task ResetConversationAsync()
    {
        CancelAnyCurrentResponse();
        messages.Clear();
        messages.Add(new(ChatRole.System, SystemPrompt));
        chatSuggestions?.Clear();
        await chatInput!.FocusAsync();
    }

    [Description("Searches for information using a phrase or keyword")]
    private async Task<IEnumerable<string>> SearchAsync(
        [Description("The phrase to search for.")] string searchPhrase,
        [Description("If possible, specify the filename to search that file only. If not provided or empty, the search includes all files.")] string? filenameFilter = null)
    {
        await InvokeAsync(StateHasChanged);
        var results = await Search.SearchAsync(searchPhrase, filenameFilter, maxResults: 5);
        return results.Select(result =>
            $"<result filename=\"{result.FileName}\" page_number=\"{result.PageNumber}\">{result.Text}</result>");
    }

    private async Task<string> GetWeather([Description("The city, correctly capitalized")] string city)
    {
        string[] weatherValues = ["Sunny", "Cloudy", "Rainy", "Snowy", "Balmy", "Bracing"];
        return city == "London" ? "Drizzle" : weatherValues[Random.Shared.Next(weatherValues.Length)];
    }

    [Description("Get current user information")]
    private Task<string> GetCurrentUserInfo()
    {
        return Task.FromResult(CurrentUser.IsAuthenticated ?
            $"UserId: {CurrentUser.Id}, Name: {CurrentUser.UserName}, Email: {CurrentUser.Email}, Roles: {string.Join(", ", CurrentUser.Roles)}" :
            "No user information available.");
    }

    public void Dispose()
        => currentResponseCancellation?.Cancel();
}
